// SimpleChatWebSocketHandler.java
package com.rent.message.handler;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Slf4j
public class SimpleChatWebSocketHandler implements WebSocketHandler {
    
    private final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket连接建立: {}", session.getId());
        // 可以在这里发送欢迎消息
        sendMessage(session, createSystemMessage("连接成功"));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String payload = message.getPayload().toString();
        log.info("收到消息: {}", payload);
        
        try {
            Map<String, Object> messageData = objectMapper.readValue(payload, Map.class);
            String type = (String) messageData.get("type");
            
            switch (type) {
                case "register":
                    handleRegister(session, messageData);
                    break;
                case "chat":
                    handleChatMessage(session, messageData);
                    break;
                case "test":
                    handleTestMessage(session);
                    break;
                default:
                    sendMessage(session, createErrorMessage("未知消息类型: " + type));
            }
        } catch (Exception e) {
            log.error("消息处理错误", e);
            sendMessage(session, createErrorMessage("消息格式错误"));
        }
    }

    private void handleRegister(WebSocketSession session, Map<String, Object> messageData) throws IOException {
        String userId = (String) messageData.get("userId");
        String username = (String) messageData.get("username");
        
        userSessions.put(userId, session);
        session.getAttributes().put("userId", userId);
        session.getAttributes().put("username", username);
        
        log.info("用户注册: {} - {}, 会话: {}", userId, username, session.getId());
        sendMessage(session, createSystemMessage("注册成功，欢迎 " + username));
    }

    private void handleChatMessage(WebSocketSession session, Map<String, Object> messageData) throws IOException {
        String fromUserId = (String) session.getAttributes().get("userId");
        String fromUsername = (String) session.getAttributes().get("username");
        String toUserId = (String) messageData.get("toUserId");
        String content = (String) messageData.get("content");
        
        if (fromUserId == null) {
            sendMessage(session, createErrorMessage("请先注册用户"));
            return;
        }
        
        log.info("用户 {} 发送消息给 {}: {}", fromUsername, toUserId, content);
        
        // 发送给接收者
        WebSocketSession targetSession = userSessions.get(toUserId);
        if (targetSession != null && targetSession.isOpen()) {
            Map<String, Object> chatMessage = new HashMap<>();
            chatMessage.put("type", "chat");
            chatMessage.put("fromUserId", fromUserId);
            chatMessage.put("fromUsername", fromUsername);
            chatMessage.put("content", content);
            chatMessage.put("timestamp", System.currentTimeMillis());
            sendMessage(targetSession, chatMessage);
            
            // 发送确认消息给发送者
            sendMessage(session, createSystemMessage("消息已发送"));
        } else {
            sendMessage(session, createSystemMessage("用户不在线，消息已保存"));
        }
    }

    private void handleTestMessage(WebSocketSession session) throws IOException {
        sendMessage(session, createSystemMessage("测试消息接收成功！"));
    }

    private void sendMessage(WebSocketSession session, Object message) throws IOException {
        if (session.isOpen()) {
            String jsonMessage = objectMapper.writeValueAsString(message);
            session.sendMessage(new TextMessage(jsonMessage));
        }
    }

    private Map<String, Object> createSystemMessage(String content) {
        Map<String, Object> message = new HashMap<>();
        message.put("type", "system");
        message.put("content", content);
        message.put("timestamp", System.currentTimeMillis());
        return message;

    }

    private Map<String, Object> createErrorMessage(String content) {
        Map<String, Object> message = new HashMap<>();
        message.put("type", "error");
        message.put("content", content);
        message.put("timestamp", System.currentTimeMillis());
        return message;
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误: {}", session.getId(), exception);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        if (userId != null) {
            userSessions.remove(userId);
            log.info("用户 {} 断开连接", userId);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}